import sys
import os

# 获取项目根目录的路径（即包含 scripts 和 modules 的目录）
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(project_root)
import time
import numpy as np
import argparse
from threading import Thread

import torch
import torch.nn.functional as F

import rclpy
import lcm
import toml
from rclpy.qos import QoSProfile, QoSHistoryPolicy
from rclpy.node import Node
from sensor_msgs.msg import LaserScan, Image
from std_msgs.msg import Float64

from utils.state_estimator_lcmt import state_estimator_lcmt
from utils.robot_control_response_lcmt import robot_control_response_lcmt
from utils.leg_control_data_lcmt import leg_control_data_lcmt
from utils.motor_ctrl_lcmt import motor_ctrl_lcmt
from utils.motor_robot_controller import Motor_Robot_Controller

from jit_scripts.jit_net import model   # 这里加载的是depth_encoder
from obs_buffer import *

msg = motor_ctrl_lcmt()
Ctrl = Motor_Robot_Controller()
Ctrl.run()
device_gpu = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 发送全0帧率的逻辑
def sendZeroFrames():
    for _ in range(8):
        msg.q_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.qd_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.kp_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.kd_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.tau_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        Ctrl.Send_cmd(msg)
        time.sleep(0.5)
        print("sending zero frame.")


class ObsSubscriber(Node):
    def __init__(self):
        super().__init__("obs_subscriber")

        self.handle_imu_thread = Thread(target=self.handle_imu)
        self.handle_contact_thread = Thread(target=self.handle_contact)
        self.handle_leg_thread = Thread(target=self.handle_leg)
        self.handle_send_cmd_thread = Thread(target=self.handle_send_cmd) # 发送到电机
        self.lcm_imu = lcm.LCM("udpm://239.255.76.67:7669?ttl=255")
        self.lcm_contact = lcm.LCM("udpm://239.255.76.67:7670?ttl=255")
        self.lcm_leg = lcm.LCM("udpm://239.255.76.67:7667?ttl=255")

        self.lcm_imu.subscribe("state_estimator", self.handle_imu_msg)
        self.lcm_contact.subscribe("robot_control_response", self.handle_contact_state)
        self.lcm_leg.subscribe("leg_control_data", self.handle_leg_state)

        qos_profile = QoSProfile(
            reliability=rclpy.qos.QoSReliabilityPolicy.BEST_EFFORT,
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=10,
        )
        self.depth_subscription = self.create_subscription(
            Image,
            "/camera/depth/image_rect_raw",
            self.depth_callback,
            qos_profile=qos_profile,
        )
        # self.subscription = self.create_subscription(
        #     Float64,
        #     '/rl_rate',
        #     self.rl_callback,
        #     10)
        # 深度图的回调函数
        self.depth_timer = self.create_timer(0.1, self.depth_model_callback)
        self.actor_timer = self.create_timer(0.02, self.rl_callback_jit)
        # self.thread_get_data = Thread(target=self.rl_callback_jit)
        # imu数据, 角速度, 加速度
        # self.omega = torch.zeros(1, 3, dtype=torch.float,device=device)
        # self.acc = torch.zeros(1, 3, dtype=torch.float,device=device)
        # 关节角度
        # self.q = torch.zeros(1, 12, dtype=torch.float,device=device)
        # 关节角速度
        # self.qd = torch.zeros(1, 12, dtype=torch.float,device=device)
        # 触地
        # self.contact = torch.zeros(1, 4, dtype=torch.float,device=device)
        # 深度图  获得的是[1, 58, 87]
        self.depth_image = torch.zeros(1, 58, 87, dtype=torch.float, device=device_gpu)
        
        self.depth_latent = torch.zeros(1, 32, dtype=torch.float,device=device)  # gpu
        # self.yaw = torch.zeros(1, 2, dtype=torch.float,device=device)
        # 电机控制消息
        self.motor_msg = motor_ctrl_lcmt()

        # depth网络模型
        self.model_depth = model(args_use.n_proprio).depth_model
        # 加载jit模型
        model_path = os.path.join(os.path.dirname(__file__), f"../models/5500-base_jit.pt")
        self.model_jit = torch.jit.load(model_path, map_location=device)
        # obs的创建
        self.obs_class = ObservationBuffer(1, args_use.n_obs)

        self.run()
        self.count = 0
        self.count_upper = 0
        print("Initialize done.")

    # 回调函数  接受数据
    def handle_imu_msg(self, channel, data):
        # 50hz
        msg = state_estimator_lcmt.decode(data)
        omega = torch.tensor(msg.omegaWorld, dtype=torch.float32, device=device).view(1, -1)
        acc = torch.tensor(msg.aWorld, dtype=torch.float32, device=device).view(1, -1)
        self.obs_class.obs_buf[:, : 3] = omega
        self.obs_class.obs_buf[:, 3: 6] = acc

    def handle_contact_state(self, channel, data):
        # 50hz
        msg = robot_control_response_lcmt().decode(data)

        binary_str = f"{msg.contact:04b}"

        # contact 中的顺序: 左后, 右后, 左前, 右前
        leg_data = [0.5 if digit == "0" else -0.5 for digit in [2, 3, 0, 1]]

        contact = torch.tensor(leg_data, dtype=torch.float32, device=device).view(1, -1)
        self.obs_class.obs_buf[:, 49: 53] = contact


    def handle_leg_state(self, channel, data):
        msg = leg_control_data_lcmt().decode(data)
        q = torch.tensor(msg.q, dtype=torch.float32, device=device).view(1, -1)
        qd = torch.tensor(msg.qd, dtype=torch.float32, device=device).view(1, -1)
        self.obs_class.obs_buf[:, 13: 25] = (q - self.obs_class.default_pos) * self.obs_class.dof_pos
        self.obs_class.obs_buf[:, 25: 37] = qd * self.obs_class.dof_vel

    # 下面四个是处理lcm
    def handle_imu(self):
        try:
            while True:
                self.lcm_imu.handle()
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def handle_contact(self):
        try:
            while True:
                self.lcm_contact.handle()
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def handle_leg(self):
        try:
            while True:
                self.lcm_leg.handle()
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def handle_send_cmd(self):
        try:
            while True:
                Ctrl.Send_cmd(self.motor_msg)
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def run(self):
        self.handle_imu_thread.start()
        self.handle_contact_thread.start()
        self.handle_leg_thread.start()
        self.handle_send_cmd_thread.start()

        # self.thread_get_data.start()


    # 这里只是深度相机的获取数据和第一步处理
    def depth_callback(self, msg):
        data = np.frombuffer(msg.data, dtype=np.uint16)
        image = np.reshape(data, (msg.height, msg.width))
        image = image.astype(np.float32)
        tensor_image = (torch.tensor(image, dtype=torch.float32, device=device_gpu).unsqueeze(0).unsqueeze(0))
        downsampled_image = F.interpolate(
            tensor_image, size=(58, 87), mode="bilinear", align_corners=False
        )

        self.depth_image = downsampled_image.view(1, 58, 87)
        # with torch.no_grad():
        #     depth_latent_and_yaw = self.model_depth(self.depth_image.to(device), self.obs_class.obs_buf[:, :args_use.n_proprio].to(device))
        #     self.depth_latent = depth_latent_and_yaw[:, :-2]  # gpu
        #     self.yaw = depth_latent_and_yaw[:, -2:] * 1.5

    # 这个函数的原来作用是使用gpu来处理深度信息获得输出,但是我怕他俩同时竞争gpu资源,先放到一个线程里测一下时间
    def depth_model_callback(self):
        # 这里只写(1, 58, 87)的深度图经过深度网络
        with torch.no_grad():
            depth_latent_and_yaw = self.model_depth(self.depth_image, self.obs_class.obs_buf[:, :args_use.n_proprio].clone().to(device_gpu))
            self.depth_latent = depth_latent_and_yaw[:, :-2].cpu()  # gpu
            yaw = (depth_latent_and_yaw[:, -2:] * 1.5).cpu()
            self.obs_class.obs_buf[:, 6: 8] = yaw
            torch.cuda.empty_cache()

    def rl_callback_jit(self):
        # count = 0
        # count_upper = 0
        # while True:
        # start = time.time()
        with torch.no_grad():                
            
            # 在这里首先更新下obs的前53个数据
            # 在这里更新53个数据很费时    将这53个数据的更新放到每个数据的获得地方   但是为什么插入这53个数据会费时呢
            # self.obs_class.insert_proprio(self.omega, self.acc, self.yaw, self.q, self.qd, self.contact)
            actions = self.model_jit(self.obs_class.obs_buf, self.depth_latent) # 在cpu上都是0.01左右的时间
            # 奇怪为什么这一步在cpu上能是0.01,但是到gpu上却到了0.02
            # 现在来说这步推理的速度是很快的,cpu不到0.01
            # 发送电机

            # print(f"test:{self.obs_class.compute_torques(actions.detach(), self.obs_class.obs_buf[:, 13: 25], self.obs_class.obs_buf[:, 25: 37])}")
            result = self.obs_class.compute_torques(torch.clamp(actions.detach(), min=-2.4, max=2.4), self.obs_class.obs_buf[:, 13: 25], self.obs_class.obs_buf[:, 25: 37])
            self.motor_msg.tau_des = result.clamp(-12, 12).numpy().astype(np.float32).tolist()[0]
            print(self.motor_msg.tau_des)
            self.obs_class.insert(actions, self.obs_class.obs_buf[:, :args_use.n_proprio]) # 这个时间也是0.01不到基本0.007
                
            # end_time = time.time()
            # count += 1
            # if end_time-start<0.02:
            #     count_upper += 1
            # print(f"比例:{count}/{count_upper}")    # cpu比例:3939/2932   gpu比例:4011/2675
            # print(f"1111:{end_time-start}")

    def destroy(self):
        rclpy.shutdown()
        self.handle_imu_thread.join()
        self.handle_contact_thread.join()
        self.handle_leg_thread.join()
        self.handle_send_cmd_thread.join()
        sendZeroFrames()

        

def main(args=None):
    sendZeroFrames()
    rclpy.init(args=args)
    obs_subscriber = ObsSubscriber()
    try:
        rclpy.spin(obs_subscriber)
    except KeyboardInterrupt:
        print("Ctrl+C pressed. Shutting down.")
    finally:
        # 清理和退出
        obs_subscriber.destroy()
        Ctrl.quit()
        sys.exit()

parser = argparse.ArgumentParser()
parser.add_argument("--n_obs", type=int, default=753)
parser.add_argument("--n_proprio", type=int, default=53)
parser.add_argument("--n_scan", type=int, default=132)
parser.add_argument("--n_actions", type=int, default=12)
parser.add_argument("--n_priv_latent", type=int, default=29)
parser.add_argument("--n_priv_explicit", type=int, default=9)
parser.add_argument("--n_hist", type=int, default=10)
args_use = parser.parse_args()
if __name__ == "__main__":
    main()

